package com.countertack.endtoendstories;

import java.io.File;
import java.io.FileOutputStream;
import java.net.InetAddress;
import java.sql.ResultSet;
import java.util.UUID;

import com.countertack.impala.ImpalaConnector;
import com.countertack.ssh.SSHManager;
import com.countertack.utils.Configuration;
import com.countertack.utils.CustomFileFilter;
import com.countertack.utils.HttpClientUtils;
import com.countertack.vmware.VMWareConnector;
import com.countertack.vmware.VMWareUtils;
import com.vmware.vim25.mo.Task;
import com.vmware.vim25.mo.VirtualMachine;

public class MalwareTest {
	static boolean isXpmachine=false;
	public static void main(String[] args) throws Exception {

		/*if (args.length <= 0) {

			System.out.println("configuration details missed !");
			return;
		}
		String configFilePath = args[0];*/

		 String configFilePath ="C:\\Users\\user\\Desktop\\SCI_Automation_Script\\config.properties";
		// "C:\\Users\\user\\Desktop\\automation\\configm.properties";
		String uuid = UUID.randomUUID().toString();
		Configuration.loadProperties(configFilePath);

		ImpalaConnector iConnector;

		iConnector = ImpalaConnector
				.createImpalaConnectorInstance(Configuration
						.getProperty("sentinelehost"));

		InetAddress address = InetAddress.getByName(Configuration
				.getProperty("sentinelehost"));
		String sentenelIp = address.getHostAddress();
		System.out.println(address.getHostAddress());
		// connect to esxi and starting virtual machines

		VMWareConnector esxiConnector = new VMWareConnector(
				Configuration.getProperty("esxi"));
		// poweredOn
		String malwareSnapShotName = Configuration
				.getProperty("rootsnapshotname");

		final VirtualMachine endPointMachine = esxiConnector
				.getVirtualMachineByname(Configuration.getProperty("vmname"));

		esxiConnector.revertToSnapShort(endPointMachine.getSnapshot()
				.getRootSnapshotList(), malwareSnapShotName);

		System.out.println("Reverting to root snap shot>>>>>"
				+ malwareSnapShotName);

		Thread.sleep(180000);

		String endpointMachineStatus = esxiConnector
				.getVirtualMachineStatus(endPointMachine);

		esxiConnector.checkAndRunVirtualMachine(endPointMachine);

		
		String osType = VMWareUtils.getVMOSArchitecture(
				Configuration.getProperty("esxi"),
				Configuration.getProperty("vmname"));
		
		if (osType != null) {

			osType = osType.toLowerCase();

			if (osType.contains("xp")) {

				isXpmachine = true;
			}
		}


		File f = new File(Configuration.getProperty("malwares"));
		File[] files = f.listFiles(new CustomFileFilter(".exe"));

		for (File malware : files) {
			try {
				esxiConnector = new VMWareConnector(
						Configuration.getProperty("esxi"));
				runMalwareTest(esxiConnector, iConnector, sentenelIp,
						malware.getPath());
			} catch (Exception e) {

			} finally {

				try {
					if (esxiConnector != null) {

						esxiConnector.closeEsxiConnection();
					}

				} catch (Exception e) {

				}
			}

		}

		// System.exit(0);

	}

	public static void runMalwareTest(VMWareConnector esxiConnector,
			ImpalaConnector iConnector, String sentenelIp, String malwaresPath)
			throws Exception {
		String endptIp = Configuration.getProperty("vmip");
		StringBuffer resultBuffer = new StringBuffer();
		resultBuffer.append("VMIP::::" + endptIp + "\n");

		VirtualMachine endPointMachine = esxiConnector
				.getVirtualMachineByname(Configuration.getProperty("vmname"));

		// win7x64-sym-2
		System.out.println("power state >>>"
				+ endPointMachine.getRuntime().powerState);

		esxiConnector.checkAndRunVirtualMachine(endPointMachine);
		Thread.sleep(240000);

		String vmUser = Configuration.getProperty("vmuser");
		String vmPassword = Configuration.getProperty("vmpassword");

		// initialising ssh to virtual machines
		// String endptIp=endPointMachine.getGuest().ipAddress;
		System.out.println("ip>>>" + endptIp);

		SSHManager endpointSsh = new SSHManager(vmUser, vmPassword, endptIp);
		endpointSsh.connect();

		// install agent

		endpointSsh.copyFile(Configuration.getProperty("driverpath"),
				"CounterTack Agent Installer Hiding.exe");

		if (isXpmachine) {
			endpointSsh
					.runCommand("cmd /c \"C:\\Documents and Settings\\"
							+ vmUser
							+ "\\CounterTack Agent Installer Hiding.exe\" /S /HOST:"
							+ sentenelIp + " /PORT:443");

		} else {

			endpointSsh
					.runCommand("cmd /c \"C:\\Users\\"
							+ vmUser
							+ "\\CounterTack Agent Installer Hiding.exe\" /S /HOST:"
							+ sentenelIp + " /PORT:443");

		}


		System.out.println("Installing agent on endpoint");
		Thread.sleep(300000);

		String eId = null;
		ResultSet endpointIdRes = iConnector
				.executeQueryAndReturnResults("select id from endpoints where ip='"
						+ endptIp
						+ "'"
						+ " and status='ON' and ip is not null order by time_started desc limit 1");

		while (endpointIdRes.next()) {
			try {
				eId = endpointIdRes.getString(1);
			} catch (Exception e) {

			}

			System.out.println("endpoint Id>>>>>>" + eId);

		}

		if (eId == null) {

			System.out.println("endpoint is not installed");

			return;

		}
		resultBuffer.append("EndpointId::::" + eId + "\n");
		resultBuffer.append("Sentinel:::" + sentenelIp + "\n");
		String driver_version = getDriverVersion(eId, iConnector);
		System.out.println("Driver Version::::" + driver_version);
		resultBuffer.append("Kernel Module Version:::" + driver_version + "\n");
		pauseEndpoint(eId);
		Thread.sleep(60000);
		File malwaref = new File(malwaresPath);
		System.out.println("malware mname>>>" + malwaref.getName());

		String malwareTestString = malwaref.getName();
		if (malwaref.getName().length() > 4) {

			malwareTestString = malwaref.getName().substring(0, 4);

		}
		endpointSsh.copyFile(malwaresPath, malwaref.getName());

		Thread.sleep(6000);
		System.out.println("malware copied success fully");

		resumeEndpoint(eId);

		Thread.sleep(60000);

		try {
			
			
			if(isXpmachine){
				endpointSsh.runCommand("cmd /c \"C:\\Documents and Settings\\" + vmUser + "\\"
						+ malwaref.getName()+"\" /quiet");
				
			}else{
				
				
				endpointSsh.runCommand("cmd /c \"C:\\Users\\" + vmUser + "\\"
						+ malwaref.getName() + "\" /quiet");
			}
			
			System.out.println("malware installed...");

		} catch (Exception e) {

		}

		boolean isbsod = false;
		String endpointMachineStatus = esxiConnector
				.getVirtualMachineStatus(endPointMachine);
		if ("Running".equalsIgnoreCase(endpointMachineStatus)) {

			String bsodRes = endpointSsh
					.returnCommandOutPut("cmd /c cd C:\\Windows && dir /b /a-d");

			if (bsodRes != null) {

				if (bsodRes.contains("MEMORY.dump")) {

					isbsod = true;

				}

			}
		}
		Thread.sleep(1500000);

		// retriving behavior id

		ResultSet brs = iConnector
				.executeQueryAndReturnResults("select hex(id) from behaviors where uri like '%"
						+ eId + "%'  order by event_start_time desc limit 1");
		String bhid = null;
		boolean isMalwareTriggered = false;
		while (brs.next()) {
			bhid = brs.getString(1);

		}
		System.out.println("behavior id>>>>" + bhid);

		resultBuffer.append("Behavior id :::::" + bhid + "\n");
		int no_of_files = 0;
		int no_of_registries = 0;
		int no_of_processes = 0;
		int notification_count = 0;
		int key_event_count = 0;
		int event_count = 0;
		int behavior_count = 0;
		int basic_event_count = 0;

		if (bhid != null) {

			ResultSet ressProcess = iConnector
					.executeQueryAndReturnResults("select target_process_name from behaviors_denorm where hex(behavior_id) = '"
							+ bhid + "' and behavior_id is not null");
			resultBuffer
					.append("<------------------Processes Names-----------> \n");

			while (ressProcess.next()) {
				String processName = ressProcess.getString(1);

				if (processName != null) {

					if (processName.contains(malwareTestString)) {

						isMalwareTriggered = true;

					}
					resultBuffer.append(processName + "\n");
					no_of_processes++;
					System.out.println("Process name>>>" + processName);

				}
			}

			ResultSet flwrrs = iConnector
					.executeQueryAndReturnResults("select target_file_path from behaviors_denorm where hex(behavior_id) = '"
							+ bhid + "' and behavior_id is not null");
			resultBuffer.append("<---------------File Names-----------> \n");
			while (flwrrs.next()) {
				String fn = flwrrs.getString(1);

				if (fn != null) {
					resultBuffer.append(fn + "\n");
					System.out.println("file name>>>" + fn);
					no_of_files++;

				}
			}

			ResultSet regwrrs = iConnector
					.executeQueryAndReturnResults("select target_registry_key from behaviors_denorm where hex(behavior_id) = '"
							+ bhid + "' and behavior_id is not null");
			// regkey
			resultBuffer.append("<-------------Registry Keys-----------> \n");
			while (regwrrs.next()) {
				String rn = regwrrs.getString(1);

				if (rn != null) {
					resultBuffer.append(rn + "\n");
					System.out.println("Reg Key>>>" + rn);
					no_of_registries++;

				}
			}

		}

		ResultSet clsrs = iConnector
				.executeQueryAndReturnResults("select classification_name,endpoint_id from classified_endpoints where endpoint_id='"
						+ eId + "'");
		boolean clsoc = false;
		String clsName = null;
		String clsGroup = "";
		while (clsrs.next()) {
			clsName = clsrs.getString(1);

			if (clsName != null) {
				System.out.println("classification_name>>>" + clsName);
				clsGroup = clsGroup + "_" + clsName;
				clsoc = true;
				// break;

			}
		}
		clsName = clsGroup;
		if (clsName.startsWith("_")) {

			clsName = clsName.replaceFirst("_", "");
		}
		if (clsoc) {

			System.out.println("Malware test result -------->Pass");
		} else {

			System.out.println("Malware test result ---------->Fail");

		}

		try {
			ResultSet countRs = iConnector
					.executeQueryAndReturnResults("select notification_count,key_event_count,event_count,behavior_count,basic_event_count from endpoints where id='"
							+ eId + "'");

			while (countRs.next()) {
				try {
					notification_count = countRs.getInt(1);
					key_event_count = countRs.getInt(2);
					event_count = countRs.getInt(3);
					behavior_count = countRs.getInt(4);
					basic_event_count = countRs.getInt(5);
				} catch (Exception e) {

					e.printStackTrace();

				}

			}

		} catch (Exception e1) {

			e1.printStackTrace();

		}

		String malName = malwaref.getName();

		malName = malName.replaceAll("\\.", "_");

		File malResultFile = new File("malware_result_"
				+ Configuration.getProperty("vmname") + "_" + malName + ".txt");

		malResultFile.createNewFile();
		try {
			FileOutputStream fout = new FileOutputStream(malResultFile);
			resultBuffer.append("Malware behavior status:" + isMalwareTriggered
					+ "\n");
			resultBuffer
					.append("Malware classification name:" + clsName + "\n");
			resultBuffer.append("BSOD:" + isbsod + "\n");

			resultBuffer.append("Number of Files:" + no_of_files + "\n");

			resultBuffer.append("Number of Regestries:" + no_of_registries
					+ "\n");

			resultBuffer
					.append("Number of Processes:" + no_of_processes + "\n");
			resultBuffer
					.append("Number of Keyevents:" + key_event_count + "\n");

			resultBuffer.append("Number of events:" + event_count + "\n");

			//

			fout.write(resultBuffer.toString().getBytes());

			fout.flush();
			fout.close();

		} catch (Exception e) {

			e.printStackTrace();

		}

		// reverting to snap shot
		uninstallEndPoint(eId);
		Thread.sleep(60000);
		while (true) {
			try {

				esxiConnector.revertToSnapShort(endPointMachine.getSnapshot()
						.getRootSnapshotList(), Configuration
						.getProperty("rootsnapshotname"));
				break;
			} catch (Exception e) {

				try {
					esxiConnector = new VMWareConnector(
							Configuration.getProperty("esxi"));
					endPointMachine = esxiConnector
							.getVirtualMachineByname(Configuration
									.getProperty("vmname"));

				} catch (Exception e1) {

					e1.printStackTrace();
				}

			}

		}
		
		try{
			endpointSsh.closeSession();
			endpointSsh.doReconnect=false;
		}catch(Exception e){
			
			
		}

		Thread.sleep(720000);

	}

	public static void pauseEndpoint(String eid) throws Exception {

		String pause = "{ " + "  \"type\":\"pause_many\", "
				+ "  \"endpoints\":[ " + "\"" + eid + "\" " + " ] " + "} ";

		HttpClientUtils.composePostRequest(
				Configuration.getProperty("restapirooturi")
						+ "/endpoints/command", pause);

	}

	public static void resumeEndpoint(String eid) throws Exception {

		String resume = "{ " + "  \"type\":\"resume_many\", "
				+ "  \"endpoints\":[ " + "    \"" + eid + "\" " + "] " + "} ";

		HttpClientUtils.composePostRequest(
				Configuration.getProperty("restapirooturi")
						+ "/endpoints/command", resume);

	}

	public static void uninstallEndPoint(String eid) throws Exception {

		{

			String uninstall = "{ " + "  \"type\":\"uninstall_many\", "
					+ "  \"endpoints\":[ " + "    \"" + eid + "\" " + "] "
					+ "} ";

			HttpClientUtils.composePostRequest(
					Configuration.getProperty("restapirooturi")
							+ "/endpoints/command", uninstall);

		}

	}

	public static void main1(String[] args) throws Exception {

		String s = "6d1b93355676e2a51e04e88b38baaea2e01134cf46d6582580e6169f73be8274.exe";

		s = s.replaceAll("\\.", "_");

		System.out.println(s);

		// File f = new File("C:\\Users\\user\\Desktop\\malware\\");
		// File[] files = f.listFiles(new CustomFileFilter(".exe"));
		//
		// for (File f2 : files) {
		//
		// String path = f2.getPath();
		//
		// System.out.println("Path>>>" + path);
		// System.out.println("Path>>>" + f2.getAbsolutePath());
		// System.out.println("Path>>>" + f2.getCanonicalPath());
		// }

	}

	public static String getDriverVersion(String eid, ImpalaConnector iConnector) {

		String driver_Version = null;
		ResultSet endpointStatRes = iConnector
				.executeQueryAndReturnResults("select driver_version from endpoints where id='"
						+ eid + "'");
		try {
			while (endpointStatRes.next()) {
				try {
					driver_Version = endpointStatRes.getString(1);
				} catch (Exception e) {

				}

				System.out.println("Driver version>>>>>>" + driver_Version);

			}
		} catch (Exception e) {

		}

		return driver_Version;

	}

	public static void main4(String[] args) throws Exception {
		String configFilePath = "C:\\Users\\user\\Desktop\\automation\\configm.properties";
		String uuid = UUID.randomUUID().toString();
		Configuration.loadProperties(configFilePath);

		File f = new File(Configuration.getProperty("malwares"));
		File[] files = f.listFiles(new CustomFileFilter(".exe"));

		for (File malware : files) {

			System.out.println(malware.getPath());

		}

	}

}
